#include <cassert>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <vector>
#include <iostream>
#include <utility>

#include <limits>

#include "sds/sfile.h"
#include "sds/sdataview.h"

#include "cluster.h"
#include "matrix.h"

#include "random.h"

using namespace std;
using namespace sds;

const int dx[] = { 1, 0, -1, 0};
const int dy[] = { 0, 1, 0, -1};
const int ds = 4;

typedef SDataView<unsigned char> CharView;
typedef SDataView<Vector> VectorView;

struct Field {
  vector<double> likelihood;
  //  vector<double> distance;
};

typedef SDataView<Field> FieldView;

const double maxdist = 80.0;

Random rng;

void Read(const string& filename, const string& refname, 
	  VectorView& view, vector<int>& id, vector<Matrix>& inverse, 
	  vector<double>& lndet, vector<Vector>& center, vector<SInteger>& size) {
  SFile file(filename);

  CharView source[CH];

  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<unsigned char>();
  }
  
  size = source[0].getSize();

  SFile reffile(refname);
  CharView ref = reffile.getDataset(0).getDataView<unsigned char>();
  
  view = VectorView(size);

  Cluster all[256];

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Vector vec(CH);
      Point point;
      for (int k = 0; k < CH; ++k) {
	vec(k) = point[k] = source[k][i][j];
      }
      if (ref[i][j] != 0) {
	all[ref[i][j]] += point;
      }
      view[i][j] = vec;
    }
  }

  for (int k = 0; k < 256; ++k) {
    if (all[k].size() > 0) {
      Matrix m(CH, CH);
      Vector v(CH);
      for (int i = 0; i < CH; ++i) {
	for (int j = 0; j < CH; ++j) {
	  m(i, j) = all[k].covariance(i, j);
	}
	v(i) = all[k].expected(i);
      }
      id.push_back(k);
      lndet.push_back(log(fabs(m.determinant())));
      inverse.push_back(m.inverse());
      center.push_back(v);
    }
  }
}


int main(int argc, const char * argv[]) {
  assert(argc == 4);

  VectorView view;
  CharView ref;

  vector<SInteger> size;
  vector<int> id;
  vector<Matrix> inverse;
  vector<Vector> center;
  vector<double> lndet;

  Read(argv[1], argv[2], view, id, inverse, lndet, center, size);

  int state = id.size();
  cout << "State: " << state << endl;

  FieldView fld(size);
  
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Field& field = fld[i][j];
      for (int l = 0; l < state; ++l) {
	Vector current = view[i][j] - center[l];
	field.likelihood.push_back(current * inverse[l] * current + lndet[l]);
      }
//       for (int k = 0; k < ds; ++k) {
// 	double dist = 0.0;
// 	if (i + dx[k] >= 0 && i + dx[k] < size[0]) {
// 	  if (j + dy[k] >= 0 && j + dy[k] < size[1]) {
// 	    dist = distance_rp((Point&)view[i][j], view[i + dx[k]][j + dy[k]]);
// 	  }
// 	}
// 	field.distance.push_back(dist);
//       }
    }
  }


  CharView curr(size), prev(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      char cl = (char)(state * rng.uniform()); 
      prev[i][j] = id[cl];
    }
  }

  const int length = 20;

  for (int k = 0; k < length; ++k) {
    cout << "Step: " << k << endl;
    double t = 3.0 / log (1.0 + k * 0.15);
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	Field& field = fld[i][j];
	vector<double> probs(state);
	double sum = 0.0;
	for (int l = 0; l < state; ++l) {
	  double dist = field.likelihood[l];
	  //	  std::cout << dist << ' ';
	  for (int m = 0; m < ds; ++m) {
	    if (i + dx[m] < 0 || i + dx[m] >= size[0]) continue;
	    if (j + dy[m] < 0 || j + dy[m] >= size[1]) continue;
	    if (prev[i + dx[m]][j + dy[m]] != id[l]) dist += 15.0;
	  }
	  //	  std::cout << dist << std::endl;
	  probs[l] = exp( - dist / t);
	  sum += probs[l];
	}
	double prob = rng.uniform() * sum;
	for (int l = 0; l < (int)state; ++l) {
	  if (prob < probs[l]) {
	    curr[i][j] = id[l];
	    break;
	  }
	  prob -= probs[l];
	}
      }
    }
    curr.swap(prev);
  } 

  SFile output(argv[3], SFile::CREATE);
  SDataset out = output.createDataset("Classified Image", UINT8, size);
  out.setDataView<unsigned char>(prev);
  return 0;
}
